#include "FaceRecognizer.h"
#include <fstream>
#include <sstream>
#include <iostream>
#include <dlib/image_processing.h>
#include <dlib/image_io.h>

FaceRecognizer::FaceRecognizer(const std::map<std::string, std::string>& config)
    : config_(config), initialized_(false) {
    threshold_ = 0.6;
    if (config_.find("threshold") != config_.end()) {
        threshold_ = std::stod(config_.at("threshold"));
    }
}

FaceRecognizer::~FaceRecognizer() {}

bool FaceRecognizer::initialize() {
    try {
        // 初始化人脸检测器
        detector_ = dlib::get_frontal_face_detector();
        
        // 加载关键点检测模型
        if (config_.find("shape_predictor_path") != config_.end()) {
            dlib::deserialize(config_.at("shape_predictor_path")) >> sp_;
        } else {
            std::cerr << "配置文件中缺少 shape_predictor_path" << std::endl;
            return false;
        }
        
        // 加载人脸识别模型
        if (config_.find("face_rec_model_path") != config_.end()) {
            dlib::deserialize(config_.at("face_rec_model_path")) >> net_;
        } else {
            std::cerr << "配置文件中缺少 face_rec_model_path" << std::endl;
            return false;
        }
        
        // 加载人脸库
        face_lib_ = loadFaceLib();
        if (face_lib_.empty()) {
            std::cerr << "人脸库为空或加载失败" << std::endl;
        }
        
        initialized_ = true;
        return true;
    } catch (const std::exception& e) {
        std::cerr << "初始化失败: " << e.what() << std::endl;
        return false;
    }
}

void FaceRecognizer::printFaceLibInfo() const {
    std::cout << "人脸库信息:" << std::endl;
    std::cout << "  人数: " << face_lib_.size() << std::endl;
    
    std::map<std::string, int> name_count;
    for (const auto& record : face_lib_) {
        name_count[record.name]++;
    }
    
    for (const auto& pair : name_count) {
        std::cout << "  " << pair.first << ": " << pair.second << " 个样本" << std::endl;
    }
    std::cout << "  匹配阈值: " << threshold_ << std::endl;
}

std::vector<dlib::rectangle> FaceRecognizer::detectFaces(const dlib::matrix<dlib::rgb_pixel>& img) {
    if (!initialized_) {
        std::cerr << "人脸识别器未初始化" << std::endl;
        return {};
    }
    return detector_(img);
}

dlib::matrix<dlib::rgb_pixel> FaceRecognizer::alignFace(const dlib::matrix<dlib::rgb_pixel>& img, 
                                                      const dlib::rectangle& face) {
    if (!initialized_) {
        throw std::runtime_error("人脸识别器未初始化");
    }
    
    auto shape = sp_(img, face);
    dlib::matrix<dlib::rgb_pixel> face_chip;
    dlib::extract_image_chip(img, dlib::get_face_chip_details(shape, 150, 0.25), face_chip);
    return face_chip;
}

std::string FaceRecognizer::recognizeFace(const dlib::matrix<dlib::rgb_pixel>& face_chip) {
    if (!initialized_) {
        return "系统未初始化";
    }
    
    dlib::matrix<float,0,1> face_descriptor = net_(face_chip);
    
    std::string best_match = "未知";
    double min_distance = threshold_ + 0.1; // 确保大于阈值
    
    for (const auto& record : face_lib_) {
        double distance = calculateDistance(face_descriptor, record.descriptor);
        if (distance < min_distance) {
            min_distance = distance;
            best_match = record.name;
        }
    }
    
    if (min_distance > threshold_) {
        return "陌生人";
    }
    
    return best_match;
}

double FaceRecognizer::compareFaces(const dlib::matrix<dlib::rgb_pixel>& face1, 
                                  const dlib::matrix<dlib::rgb_pixel>& face2) {
    if (!initialized_) {
        return -1.0;
    }
    
    dlib::matrix<float,0,1> desc1 = net_(face1);
    dlib::matrix<float,0,1> desc2 = net_(face2);
    return calculateDistance(desc1, desc2);
}

void FaceRecognizer::saveFaceDescriptor(const std::string& name, 
                                      const dlib::matrix<dlib::rgb_pixel>& face_chip) {
    if (!initialized_) {
        return;
    }
    
    dlib::matrix<float,0,1> descriptor = net_(face_chip);
    
    // 添加到内存中的人脸库
    face_lib_.push_back({name, descriptor});
    
    // 更新CSV文件 - 修改键名为 face_library
    if (config_.find("face_library") != config_.end()) {  // 修改这里
        std::ofstream file(config_.at("face_library"), std::ios::app);  // 修改这里
        if (file) {
            file << name;
            for (long i = 0; i < descriptor.size(); ++i) {
                file << "," << descriptor(i);
            }
            file << "\n";
        }
    }
}

void FaceRecognizer::updateFaceLib() {
    face_lib_ = loadFaceLib();
}

std::vector<FaceRecognizer::FaceRecord> FaceRecognizer::loadFaceLib() {
    std::vector<FaceRecord> lib;
    
    // 修改键名为 face_library
    if (config_.find("face_library") == config_.end()) {  // 修改这里
        std::cerr << "配置文件中缺少 face_library 路径" << std::endl;
        return lib;
    }
    
    std::ifstream file(config_.at("face_library"));  // 修改这里
    if (!file) {
        std::cerr << "无法打开人脸库文件: " << config_.at("face_library") << std::endl;
        return lib;
    }
    
    std::string line;
    while (std::getline(file, line)) {
        std::istringstream iss(line);
        FaceRecord record;
        std::string value;
        
        std::getline(iss, record.name, ',');
        
        dlib::matrix<float,0,1> descriptor(128);
        for (long i = 0; i < 128; ++i) {
            if (!std::getline(iss, value, ',')) break;
            try {
                descriptor(i) = std::stof(value);
            } catch (...) {
                break;
            }
        }
        record.descriptor = descriptor;
        lib.push_back(record);
    }
    
    std::cout << "已从 " << config_.at("face_library") 
              << " 加载 " << lib.size() << " 个人脸特征" << std::endl;
    return lib;
}

double FaceRecognizer::calculateDistance(const dlib::matrix<float,0,1>& A, 
                                       const dlib::matrix<float,0,1>& B) const {
    return dlib::length(A - B);
}

double FaceRecognizer::calculateCosineSimilarity(const dlib::matrix<float,0,1>& A, 
                                               const dlib::matrix<float,0,1>& B) const {
    double dot_product = dlib::dot(A, B);
    double norm_a = dlib::length(A);
    double norm_b = dlib::length(B);
    return dot_product / (norm_a * norm_b);
}